import torch
import torch.nn as nn
from math import sqrt
from basicsr.utils.registry import ARCH_REGISTRY


class Conv_ReLU_Block(nn.Module):
    """原始VDSR论文中的卷积-ReLU块"""
    def __init__(self, num_feat):
        super(Conv_ReLU_Block, self).__init__()
        self.conv = nn.Conv2d(in_channels=num_feat, out_channels=num_feat, kernel_size=3, stride=1, padding=1, bias=False)
        self.relu = nn.ReLU(inplace=True)

    def forward(self, x):
        return self.relu(self.conv(x))


@ARCH_REGISTRY.register()
class VDSR(nn.Module):
    """VDSR模型实现，严格按照原始论文设计，支持RGB三通道输入

    Paper: Accurate Image Super-Resolution Using Very Deep Convolutional Networks

    Args:
        num_in_ch (int): 输入通道数. Default: 3 (RGB图像).
        num_out_ch (int): 输出通道数. Default: 3 (RGB图像).
        num_feat (int): 中间特征通道数. Default: 64.
        upscale (int): 上采样因子. 支持2^n和3. Default: 2.
    """

    def __init__(self,
                 num_in_ch=3,
                 num_out_ch=3,
                 num_feat=64,
                 upscale=2):
        super(VDSR, self).__init__()

        self.upscale = upscale
        self.num_in_ch = num_in_ch
        self.num_out_ch = num_out_ch

        # 输入卷积层 - 将输入映射到特征空间
        self.input = nn.Conv2d(in_channels=num_in_ch, out_channels=num_feat, kernel_size=3, stride=1, padding=1, bias=False)

        # 18个残差卷积层
        self.residual_layer = self.make_layer(Conv_ReLU_Block, 18, num_feat)

        # 输出卷积层 - 将特征映射回图像空间
        self.output = nn.Conv2d(in_channels=num_feat, out_channels=num_out_ch, kernel_size=3, stride=1, padding=1, bias=False)

        # ReLU激活函数
        self.relu = nn.ReLU(inplace=True)

        # 上采样层
        self.upsample = nn.Upsample(scale_factor=upscale, mode='bicubic', align_corners=False)

        # 权重初始化 - 按照原始论文的方式
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0, sqrt(2. / n))

    def make_layer(self, block, num_of_layer, num_feat):
        """创建指定数量的卷积层"""
        layers = []
        for _ in range(num_of_layer):
            layers.append(block(num_feat))
        return nn.Sequential(*layers)

    def forward(self, x):
        # 上采样输入到目标分辨率
        if self.upscale > 1:
            x = self.upsample(x)

        # 存储残差（用于残差学习）
        residual = x

        # 前向传播
        out = self.relu(self.input(x))  # 输入卷积+激活
        out = self.residual_layer(out)   # 18个残差块
        out = self.output(out)           # 输出卷积

        # 残差学习 - 将学习到的残差添加到输入
        out = torch.add(out, residual)

        return out